package com.gitee.ixtf.mongo.kotlinx

import com.gitee.ixtf.core.EntityDTO
import com.gitee.ixtf.core.Sort
import com.gitee.ixtf.core.SortStart
import com.gitee.ixtf.mongo.Jmongo.Companion.DEFAULT_BATCH_SIZE
import com.gitee.ixtf.mongo.Jmongo.Companion.idEq
import com.gitee.ixtf.mongo.JmongoRef
import com.mongodb.client.model.Filters.and
import com.mongodb.client.model.IndexOptions
import com.mongodb.client.model.Sorts
import com.mongodb.kotlin.client.coroutine.FindFlow
import com.mongodb.kotlin.client.coroutine.MongoCollection
import com.mongodb.kotlin.client.model.Indexes
import java.security.Principal
import kotlinx.coroutines.flow.*
import kotlinx.coroutines.runBlocking
import org.bson.BsonDocument
import org.bson.BsonDocumentReader
import org.bson.BsonDocumentWriter
import org.bson.Document
import org.bson.codecs.DecoderContext
import org.bson.codecs.EncoderContext
import org.bson.conversions.Bson
import org.bson.types.ObjectId

fun MongoCollection<*>.ensureIndex(keyMap: Map<String, Int>, options: IndexOptions) = runBlocking {
  listIndexes().firstOrNull { it.getString("name") == options.name }
      ?: createIndex(Document(keyMap), options.background(true))
}

fun MongoCollection<*>.ensureIndex(bson: Bson, options: IndexOptions) = runBlocking {
  listIndexes().firstOrNull { it.getString("name") == options.name }
      ?: createIndex(bson, options.background(true))
}

suspend fun <T : Any> MongoCollection<T>.singleOrNull(id: ObjectId): T? =
    find(idEq(id)).singleOrNull()

suspend fun <T : Any> MongoCollection<T>.single(id: ObjectId): T = find(idEq(id)).single()

suspend fun MongoCollection<*>.exists(id: ObjectId): Boolean = countDocuments(idEq(id)) > 0

suspend fun <T : Any> MongoCollection<T>.singleOrNull(id: String): T? =
    find(idEq(id)).singleOrNull()

suspend fun <T : Any> MongoCollection<T>.single(id: String): T = find(idEq(id)).single()

suspend fun MongoCollection<*>.exists(id: String): Boolean = countDocuments(idEq(id)) > 0

suspend fun <T : Any> MongoCollection<T>.singleOrNull(id: JmongoRef<*>): T? =
    find(idEq(id)).singleOrNull()

suspend fun <T : Any> MongoCollection<T>.single(id: JmongoRef<*>): T = find(idEq(id)).single()

suspend fun MongoCollection<*>.exists(id: JmongoRef<*>): Boolean = countDocuments(idEq(id)) > 0

suspend fun <T : Any> MongoCollection<T>.singleOrNull(id: EntityDTO): T? =
    find(idEq(id)).singleOrNull()

suspend fun <T : Any> MongoCollection<T>.single(id: EntityDTO): T = find(idEq(id)).single()

suspend fun MongoCollection<*>.exists(id: EntityDTO): Boolean = countDocuments(idEq(id)) > 0

suspend fun <T : Any> MongoCollection<T>.singleOrNull(id: Principal): T? =
    find(idEq(id)).singleOrNull()

suspend fun <T : Any> MongoCollection<T>.single(id: Principal): T = find(idEq(id)).single()

suspend fun MongoCollection<*>.exists(id: Principal): Boolean = countDocuments(idEq(id)) > 0

fun <T : Any> MongoCollection<*>.asBson(o: T): BsonDocument =
    BsonDocument().also { bsonDocument ->
      BsonDocumentWriter(bsonDocument).use { writer ->
        val codec = codecRegistry.get(o.javaClass)
        codec.encode(writer, o, EncoderContext.builder().build())
      }
    }

fun <T : Any> MongoCollection<*>.fromBson(bson: BsonDocument, clazz: Class<T>): T =
    BsonDocumentReader(bson).use { reader ->
      val codec = codecRegistry.get(clazz)
      return codec.decode(reader, DecoderContext.builder().build())
    }

fun <T : Any> MongoCollection<*>.fromBson(bson: String, clazz: Class<T>): T =
    fromBson(BsonDocument.parse(bson), clazz)

inline fun <reified T : Any> MongoCollection<*>.fromBson(bson: BsonDocument): T =
    fromBson(bson, T::class.java)

inline fun <reified T : Any> MongoCollection<*>.fromBson(bson: String): T =
    fromBson(BsonDocument.parse(bson))

suspend fun MongoCollection<*>.count(condition: Bson? = null): Long =
    condition?.let { countDocuments(it) } ?: countDocuments()

suspend fun MongoCollection<*>.count(condition: Flow<Bson?>) = count(condition.asFilter())

suspend fun MongoCollection<*>.count(vararg condition: Bson?) = count(condition.asFlow())

suspend fun MongoCollection<*>.count(condition: Iterable<Bson?>) = count(condition.asFlow())

fun <T : Any> MongoCollection<T>.list(condition: Bson? = null): FindFlow<T> {
  val findFlow = condition?.let { find(it) } ?: find()
  return findFlow.allowDiskUse(true).batchSize(DEFAULT_BATCH_SIZE)
}

suspend fun <T : Any> MongoCollection<T>.list(condition: Flow<Bson?>) = list(condition.asFilter())

suspend fun <T : Any> MongoCollection<T>.list(vararg condition: Bson?) = list(condition.asFlow())

suspend fun <T : Any> MongoCollection<T>.list(condition: Iterable<Bson?>) = list(condition.asFlow())

suspend fun <T : Any> FindFlow<T>.sort(sort: Flow<Sort>?): FindFlow<T> = sort(sort?.asSort())

suspend fun <T : Any> FindFlow<T>.sort(vararg sort: Sort) = sort(sort.asFlow())

suspend fun <T : Any> FindFlow<T>.sort(sort: Iterable<Sort>?) = sort(sort?.asFlow())

suspend fun Flow<Bson?>.asFilter(): Bson? =
    filterNotNull().toList().let {
      when (it.size) {
        0 -> null
        1 -> it.first()
        else -> and(it)
      }
    }

suspend fun Flow<Sort>.asSort(): Bson? =
    filterNotNull()
        .distinctUntilChangedBy { it.id }
        .toList()
        .map { if (it.start == SortStart.asc) Sorts.ascending(it.id) else Sorts.descending(it.id) }
        .let {
          when (it.size) {
            0 -> null
            1 -> it.first()
            else -> Sorts.orderBy(it)
          }
        }
